#!/usr/bin/env bash

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

# A utility script to install the correct packages and run the tests.

set -euo pipefail

# Check if the script was called with exactly 1 argument
if [[ ${#} -ne 1 ]]; then
  echo "Error: This script requires exactly 1 argument. You provided ${#}"
  exit 1
fi
if [[ "${1}" != "bindings" && "${1}" != "core" && "${1}" != "pathfinder" ]]; then
  echo "Error: Invalid test module '${1}'. Must be 'bindings', 'core', or 'pathfinder'"
  exit 1
fi

test_module=${1}

# Unconditionally install pathfinder wheel
# (it is a direct dependency of bindings, and a transitive dependency of core)
pushd ./cuda_pathfinder
echo "Installing pathfinder wheel"
pip install ./*.whl --group test
popd

if [[ "${test_module}" == "pathfinder" ]]; then
  pushd ./cuda_pathfinder
  echo "Running pathfinder tests with " \
      "LD:${CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS} " \
      "FH:${CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS}"
  pytest -ra -s -v tests/ |& tee /tmp/pathfinder_test_log.txt
  # Fail if no "INFO test_" lines are found; capture line count otherwise
  line_count=$(grep '^INFO test_' /tmp/pathfinder_test_log.txt | wc -l)
  echo "Number of \"INFO test_\" lines: $line_count"
  popd
elif [[ "${test_module}" == "bindings" ]]; then
  echo "Installing bindings wheel"
  pushd ./cuda_bindings
  if [[ "${LOCAL_CTK}" == 1 ]]; then
    pip install "${CUDA_BINDINGS_ARTIFACTS_DIR}"/*.whl --group test
  else
    pip install $(ls "${CUDA_BINDINGS_ARTIFACTS_DIR}"/*.whl)[all] --group test
  fi
  echo "Running bindings tests"
  ${SANITIZER_CMD} pytest -rxXs -v tests/
  if [[ "${SKIP_CYTHON_TEST}" == 0 ]]; then
    ${SANITIZER_CMD} pytest -rxXs -v tests/cython
  fi
  popd
elif [[ "${test_module}" == "core" ]]; then
  # If build/test majors match: cuda.bindings is installed in the previous step.
  # If mismatch: cuda.bindings is installed from the backport branch.
  if [[ "${SKIP_CUDA_BINDINGS_TEST}" == 1 ]]; then
    echo "Installing bindings wheel"
    if [[ "${LOCAL_CTK}" == 1 ]]; then
      pip install "${CUDA_BINDINGS_ARTIFACTS_DIR}"/*.whl
    else
      pip install $(ls "${CUDA_BINDINGS_ARTIFACTS_DIR}"/*.whl)[all]
    fi
  fi
  TEST_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${CUDA_VER})"
  echo "Installing core wheel"

  FREE_THREADING=""
  if python -c 'import sys; assert not sys._is_gil_enabled()' 2> /dev/null; then
    FREE_THREADING+="-ft"
  fi

  pushd ./cuda_core
  if [[ "${LOCAL_CTK}" == 1 ]]; then
    # We already installed cuda-bindings, and all CTK components exist locally,
    # so just install the test dependencies.
    pip install "${CUDA_CORE_ARTIFACTS_DIR}"/*.whl --group "test-cu${TEST_CUDA_MAJOR}${FREE_THREADING}"
  else
    pip install $(ls "${CUDA_CORE_ARTIFACTS_DIR}"/*.whl)["cu${TEST_CUDA_MAJOR}"] --group "test-cu${TEST_CUDA_MAJOR}${FREE_THREADING}"
  fi
  echo "Running core tests"
  ${SANITIZER_CMD} pytest -rxXs -v tests/
  # Currently our CI always installs the latest bindings (from either major version).
  # This is not compatible with the test requirements.
  if [[ "${SKIP_CYTHON_TEST}" == 0 ]]; then
    ${SANITIZER_CMD} pytest -rxXs -v tests/cython
  fi
  popd
fi
